package top.jfunc.common.db.data;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.time.Duration;
import java.time.Instant;
import java.util.ArrayList;
import java.util.List;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.function.IntConsumer;

import javax.sql.DataSource;

import top.jfunc.common.utils.IoUtil;

/**
 * 高性能数据批量更新工具，支持事物
 *
 * @author <a href="https://github.com/TAKETODAY">Harry Yang</a>
 * @since 1.0 2022/11/12 16:03
 */
public abstract class DataBatchUpdater {
  private static final Logger logger = LoggerFactory.getLogger(DataBatchUpdater.class);

  private static final int DEFAULT_BATCH_SIZE = 5000;

  protected final String updateSQL;
  protected final DataSource dataSource;

  protected final boolean autoCommit;
  protected final int maxBatchRecords;

  // @Nullable
  protected final Consumer<Throwable> exceptionConsumer;

  // @Nullable
  protected final List<Function<String, String>> sqlModifiers;

  protected final IntConsumer successCallback;

  // 保存的行数
  protected int updateRows;
  protected Instant startPoint;

  protected DataBatchUpdater(
      String updateSQL, DataSource dataSource,
      boolean autoCommit, int maxBatchRecords,
      Consumer<Throwable> exceptionConsumer,
      List<Function<String, String>> sqlModifiers,
      IntConsumer successCallback) {
    this.updateSQL = updateSQL;
    this.dataSource = dataSource;
    this.autoCommit = autoCommit;
    this.sqlModifiers = sqlModifiers;
    this.maxBatchRecords = maxBatchRecords;
    this.exceptionConsumer = exceptionConsumer;
    this.successCallback = successCallback;
  }

  public String getUpdateSQL() {
    return updateSQL;
  }

  public DataSource getDataSource() {
    return dataSource;
  }

  public int getUpdateRows() {
    return updateRows;
  }

  public boolean isAutoCommit() {
    return autoCommit;
  }

  protected PreparedStatement prepareStatement(Connection con, String sql) throws SQLException {
    return con.prepareStatement(sql);
  }

  /**
   * 执行一次更新
   * <p>
   * 该方法会创建一次 JDBC 连接，
   * 该方法结束后，如果有事物会自动提交一次事物
   */
  protected final void executeUpdate(BatchCallback batchCallback) {
    startPoint = Instant.now();
    Connection connection = null;
    PreparedStatement statement = null;

    try {
      String updateSQL = getUpdateSQL();
      if (sqlModifiers != null) {
        for (Function<String, String> sqlModifier : sqlModifiers) {
          updateSQL = sqlModifier.apply(updateSQL);
        }
      }

      connection = dataSource.getConnection();
      connection.setAutoCommit(autoCommit);
      connection.setTransactionIsolation(Connection.TRANSACTION_READ_UNCOMMITTED);
      statement = prepareStatement(connection, updateSQL);

      int rows = batchCallback.accept(statement);
      if (rows > 0) {
        executeBatch(statement);
      }
      if (!autoCommit) {
        logger.info("提交更改");
        connection.commit();
      }
      if (successCallback != null) {
        successCallback.accept(rows);
      }
    }
    catch (Throwable e) {
      if (connection != null) {
        try {
          connection.rollback();
        }
        catch (SQLException ex) {
          e.addSuppressed(ex);
        }
      }
      if (exceptionConsumer != null) {
        exceptionConsumer.accept(e);
      }
      else {
        throw new IllegalStateException("数据操作失败", e);
      }
    }
    finally {
      IoUtil.close(statement);
      IoUtil.close(connection);
    }
  }

  interface BatchCallback {

    int accept(PreparedStatement statement) throws Exception;
  }

  /**
   * 隐式或显式执行更新
   * <p>
   * 保证所有数据都更新必须要在最后手动（显式）执行一次
   *
   * @param statement 要更新的语句
   * @throws SQLException 执行更新报的错
   */
  final void executeBatch(PreparedStatement statement) throws SQLException {
    int[] updateCounts = statement.executeBatch();
    this.updateRows += updateCounts.length;
    if (logger.isDebugEnabled()) {
      Instant now = Instant.now();
      Duration between = Duration.between(startPoint, now);
      logger.debug("此次更新了'{}'行数据, 总共更新了'{}'行数据, 耗时 {}",
          updateCounts.length, updateRows, between);
    }
  }

  /**
   * 中间转换 直接处理 给定的 Item list
   */
  public static <T> ProcessItemBuilder<T> processItem() {
    return new ProcessItemBuilder<>();
  }

  /**
   * 中间转换 ResultSet 到 Bean 在将 Bean 映射数据到 PreparedStatement
   */
  public static <T> DataExchangeBuilder<T> dataExchange() {
    return new DataExchangeBuilder<>();
  }

  /**
   * 直接从 ResultSet 映射数据到 PreparedStatement
   */
  public static ExtractDataBuilder directExtract() {
    return new ExtractDataBuilder();
  }

  public static <B extends Builder<B, U>, U extends DataBatchUpdater> B from(Class<B> builderClass) {
    return ExceptionUtils.sneakyThrow(builderClass::newInstance);
  }

  // Builder

  public static abstract class Builder<B extends Builder<B, U>, U extends DataBatchUpdater> {

    int batchSize = DEFAULT_BATCH_SIZE;
    boolean autoCommit = false;

    DataSource dataSource;
    String updateSQL;

    Consumer<Throwable> exceptionConsumer;
    final ArrayList<Function<String, String>> sqlModifiers = new ArrayList<>();

    IntConsumer successCallback;

    public B autoCommit(boolean autoCommit) {
      this.autoCommit = autoCommit;
      return self();
    }

    public B batchSize(int batchSize) {
      this.batchSize = batchSize;
      return self();
    }

    public B updateSQL(String updateSQL) {
      this.updateSQL = updateSQL;
      return self();
    }

    public B dataSource(DataSource dataSource) {
      this.dataSource = dataSource;
      return self();
    }

    /**
     * 可以设置多个修改器
     */
    public B modifySQL(Function<String, String> sqlModifier) {
      this.sqlModifiers.add(sqlModifier);
      return self();
    }

    public B errorHandling(Consumer<Throwable> consumer) {
      this.exceptionConsumer = consumer;
      return self();
    }

    public B onSuccess(IntConsumer successCallback) {
      this.successCallback = successCallback;
      return self();
    }

    /**
     * 配置异常处理为抛出
     */
    public B sneakyThrow() {
      return errorHandling(ExceptionUtils::sneakyThrow);
    }

    protected abstract B self();

    // build

    public U build() {
      assertState();
      return buildInternal(updateSQL, dataSource, autoCommit,
          batchSize, exceptionConsumer, sqlModifiers, successCallback);
    }

    protected final void assertState() {
    }

    protected abstract U buildInternal(String updateSQL,
        DataSource dataSource, boolean autoCommit, int batchSize,
        Consumer<Throwable> exceptionConsumer, List<Function<String, String>> sqlModifiers, IntConsumer successCallback);

  }

  public static class ProcessItemBuilder<I> extends Builder<ProcessItemBuilder<I>, ProcessItem<I>> {
    BatchStatementCallback<I> itemStatementCallback;

    @Override
    protected ProcessItemBuilder<I> self() {
      return this;
    }

    /**
     * 条目处理器，处理Bean 将Bean映射到 PreparedStatement
     */
    public ProcessItemBuilder<I> processItem(BatchStatementCallback<I> callback) {
      this.itemStatementCallback = callback;
      return this;
    }

    @Override
    protected ProcessItem<I> buildInternal(String updateSQL,
        DataSource dataSource, boolean autoCommit, int batchSize,
        Consumer<Throwable> exceptionConsumer,
        List<Function<String, String>> sqlModifiers, IntConsumer successCallback) {
      return new ProcessItem<>(updateSQL, dataSource, autoCommit,
          batchSize, exceptionConsumer, sqlModifiers, itemStatementCallback, successCallback);
    }

    public ProcessItem<I> update(List<I> items) {
      ProcessItem<I> updater = build();
      updater.update(items);
      return updater;
    }

  }

  public static class ProcessItem<T> extends DataBatchUpdater {
    final BatchStatementCallback<T> itemStatementCallback;

    ProcessItem(String updateSQL, DataSource dataSource,
        boolean autoCommit, int maxBatchRecords, Consumer<Throwable> exceptionConsumer,
        List<Function<String, String>> sqlModifiers, BatchStatementCallback<T> itemStatementCallback,
        IntConsumer successCallback) {
      super(updateSQL, dataSource, autoCommit, maxBatchRecords, exceptionConsumer, sqlModifiers, successCallback);
      this.itemStatementCallback = itemStatementCallback;
    }

    /**
     * 将输入的条目列表 通过 一系列映射后更新到对应数据库
     * <p>
     * 该方法在将大量数据更新到数据库过程中会阻塞，即使用到了多线程更新
     *
     * @param items 要更新的对象
     */
    public void update(List<T> items) {
      executeUpdate(statement -> doUpdate(statement, items, maxBatchRecords));
    }

    /**
     * 添加到 {@link PreparedStatement#addBatch()}
     * <p>
     * 默认是单线程
     *
     * @param items 要执行更新的对象
     * @param maxBatchRecords 批量更新的大小，默认添加的条数到达该值的时候会执行保存
     * @throws SQLException 更新过程中发生的异常
     * @see ProcessItem
     */
    protected int doUpdate(PreparedStatement statement, List<T> items, int maxBatchRecords) throws Exception {
      int currentBatchRecords = 0;
      for (T item : items) {
        itemStatementCallback.applyTo(item, statement);
        statement.addBatch();

        // 判断是否隐式保存
        if (++currentBatchRecords % maxBatchRecords == 0) {
          executeBatch(statement);
          currentBatchRecords = 0;
        }
      }
      return items.size();
    }

  }

  public static class DataExchangeBuilder<I> extends Builder<DataExchangeBuilder<I>, DataExchange<I>> {
    private ResultSetExtractor<I> itemMapper;
    private BatchStatementCallback<I> itemStatementCallback;

    @Override
    protected DataExchangeBuilder<I> self() {
      return this;
    }

    /**
     * 条目处理器，处理Bean 将Bean映射到 PreparedStatement
     */
    public DataExchangeBuilder<I> processItem(BatchStatementCallback<I> callback) {
      this.itemStatementCallback = callback;
      return this;
    }

    /**
     * 条目处理器，处理Bean 将Bean映射到 PreparedStatement
     */
    public DataExchangeBuilder<I> map(ResultSetExtractor<I> itemMapper) {
      this.itemMapper = itemMapper;
      return this;
    }

    @Override
    protected DataExchange<I> buildInternal(String updateSQL,
        DataSource dataSource, boolean autoCommit, int batchSize,
        Consumer<Throwable> exceptionConsumer,
        List<Function<String, String>> sqlModifiers, IntConsumer successCallback) {

      return new DataExchange<>(updateSQL, dataSource, autoCommit,
          batchSize, exceptionConsumer, sqlModifiers, itemStatementCallback, successCallback);
    }

    /**
     * 从一个输入源开始执行数据传输和保存
     *
     * @param extractor 数据源
     */
    public DataExchange<I> from(DataExtractor extractor) {
      DataExchange<I> updater = build();
      updater.transfer(extractor, itemMapper);
      return updater;
    }

    public DataExchange<I> from(Consumer<DataExtractor.PageableBuilder> builder) {
      DataExtractor.PageableBuilder pageReading = DataExtractor.pageReading();
      builder.accept(pageReading);
      DataExtractor extractor = pageReading.build();
      return from(extractor);
    }

    public DataExchange<I> fromSimple(Consumer<DataExtractor.SimpleBuilder> builder) {
      DataExtractor.SimpleBuilder pageReading = DataExtractor.simple();
      builder.accept(pageReading);
      DataExtractor extractor = pageReading.build();
      return from(extractor);
    }

  }

  // 数据交换 ，中间有个 items 列表维护了 T

  public static class DataExchange<T> extends ProcessItem<T> {

    DataExchange(String updateSQL, DataSource dataSource, boolean autoCommit,
        int maxBatchRecords, Consumer<Throwable> exceptionConsumer,
        List<Function<String, String>> sqlModifiers,
        BatchStatementCallback<T> itemStatementCallback, IntConsumer successCallback) {
      super(updateSQL, dataSource, autoCommit, maxBatchRecords, exceptionConsumer, sqlModifiers, itemStatementCallback, successCallback);
    }

    public void transfer(DataExtractor extractor, ResultSetExtractor<T> itemMapper) {
      class ItemRowReader implements RowCallbackHandler {
        private final PreparedStatement statement;
        final ArrayList<T> items = new ArrayList<>();

        ItemRowReader(PreparedStatement statement) {
          this.statement = statement;
        }

        @Override
        public void processRow(ResultSet rs, int rowNumber) throws SQLException {
          T item = itemMapper.extractData(rs);
          items.add(item);
        }

        @Override
        public void finishedReading(int rows) throws Exception {
          doUpdate(statement, items, maxBatchRecords);
          items.clear();
        }

      }

      executeUpdate(statement -> {
        RowCallbackHandler rowHandler = new ItemRowReader(statement);
        return extractor.processRow(rowHandler);
      });
    }

  }

  public static class ExtractDataBuilder extends Builder<ExtractDataBuilder, ExtractResultSet> {
    BatchStatementCallback<ResultSet> statementCallback;

    @Override
    protected ExtractDataBuilder self() {
      return this;
    }

    /**
     * 条目处理器，处理ResultSet 将 ResultSet 映射到 PreparedStatement
     */
    public ExtractDataBuilder map(BatchStatementCallback<ResultSet> statementCallback) {
      this.statementCallback = statementCallback;
      return this;
    }

		/**
		 * 直接将 ResultSet 映射到 PreparedStatement
		 *
		 * @param columns 列数
		 */
    public ExtractDataBuilder mapColumns(int columns) {
      this.statementCallback = BatchStatementCallback.forColumns(columns);
      return this;
    }

    /**
     * 直接将 ResultSet 映射到 PreparedStatement
     *
     * @param columns 列数
     */
    public ExtractDataBuilder mapColumns(int start, int columns) {
      this.statementCallback = BatchStatementCallback.forColumns(start, columns);
      return this;
    }

		/**
		 * 直接将 ResultSet 映射到 PreparedStatement
		 */
		public ExtractDataBuilder autoColumns() {
			this.statementCallback = BatchStatementCallback.autoColumns();
			return this;
		}

    /**
     * 直接将 ResultSet 映射到 PreparedStatement
     * <p>
     * 映射之后执行的操作
     */
    public ExtractDataBuilder mapLast(BatchStatementCallback<ResultSet> lastStatementCallback) {
      BatchStatementCallback<ResultSet> firstStatementCallback = this.statementCallback;
      if (firstStatementCallback != null) {
        this.statementCallback = (item, statement) -> {
          firstStatementCallback.applyTo(item, statement);
          lastStatementCallback.applyTo(item, statement);
        };
      }
      else {
        this.statementCallback = lastStatementCallback;
      }
      return this;
    }

    @Override
    protected ExtractResultSet buildInternal(String updateSQL,
        DataSource dataSource, boolean autoCommit, int batchSize,
        Consumer<Throwable> exceptionConsumer,
        List<Function<String, String>> sqlModifiers, IntConsumer successCallback) {

      return new ExtractResultSet(updateSQL, dataSource, autoCommit,
          batchSize, exceptionConsumer, sqlModifiers, statementCallback, successCallback);
    }

    /**
     * 从一个输入源开始执行数据传输和保存
     * <p>
     * 异常传递问题，当读取发生异常了，读取处理了异常没有向上抛出，收集的数据不完整
     *
     * @param extractor 数据源
     */
    public ExtractResultSet from(DataExtractor extractor) {
      ExtractResultSet updater = build();
      updater.transferFrom(extractor);
      return updater;
    }

    /**
     * 异常传递问题，当读取发生异常了，读取处理了异常没有向上抛出，收集的数据不完整
     */
    public <T extends DataExtractor.Builder<T>> ExtractResultSet from(DataExtractor.Builder<T> builder) {
      return from(builder.build());
    }

    /**
     * 异常传递问题，当读取发生异常了，读取处理了异常没有向上抛出，收集的数据不完整
     */
    public ExtractResultSet from(Consumer<DataExtractor.PageableBuilder> builder) {
      DataExtractor.PageableBuilder pageReading = DataExtractor.pageReading();
      builder.accept(pageReading);
      return from(pageReading);
    }

    /**
     * 异常传递问题，当读取发生异常了，读取处理了异常没有向上抛出，收集的数据不完整
     */
    public ExtractResultSet fromSimple(Consumer<DataExtractor.SimpleBuilder> builder) {
      DataExtractor.SimpleBuilder pageReading = DataExtractor.simple();
      builder.accept(pageReading);
      return from(pageReading);
    }

  }

  public static class ExtractResultSet extends DataBatchUpdater {
    private final BatchStatementCallback<ResultSet> statementCallback;

    private ExtractResultSet(String updateSQL, DataSource dataSource,
        boolean autoCommit, int maxBatchRecords, Consumer<Throwable> exceptionConsumer,
        List<Function<String, String>> sqlModifiers, BatchStatementCallback<ResultSet> statementCallback, IntConsumer successCallback) {
      super(updateSQL, dataSource, autoCommit, maxBatchRecords, exceptionConsumer, sqlModifiers, successCallback);
      this.statementCallback = statementCallback;
    }

    public void transferFrom(DataExtractor extractor) {
      executeUpdate(statement ->
          extractor.processRow(new AddBatchRowHandler(statement, statementCallback)));
    }

    class AddBatchRowHandler implements RowCallbackHandler {
      private final PreparedStatement statement;
      private final BatchStatementCallback<ResultSet> statementCallback;

      private int currentBatchRecords = 0;

      AddBatchRowHandler(PreparedStatement statement,
          BatchStatementCallback<ResultSet> statementCallback) {
        this.statement = statement;
        this.statementCallback = statementCallback;
      }

      @Override
      public void processRow(ResultSet rs, int rowNumber) throws Exception {
        statementCallback.applyTo(rs, statement);
        statement.addBatch();

        // 判断是否隐式保存
        if (++currentBatchRecords % maxBatchRecords == 0) {
          executeBatch(statement);
          currentBatchRecords = 0;
        }
      }

      @Override
      public void finishedReading(int rows) throws SQLException {
        if (currentBatchRecords > 0) {
          executeBatch(statement);
        }
      }

    }

  }

}
